import os

import cv2
from ultralytics import YOLO

from exchange_opt import *

weights = r"/home/champrin/Downloads/best (1).pt"

# Load a model
model = YOLO(weights)  # pretrained YOLO11n model

slot_opt = SlotDetector(0.05, 40, 600)


image_path = "/home/champrin/Desktop/MV-CS016-10UC+DA1041860/给预备役/side/blue/b1/image"








# 获取文件夹A中所有图片文件路径
image_files = [os.path.join(image_path, f) for f in os.listdir(image_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif'))]

# 根据文件名进行排序
sorted_image_files = sorted(image_files, key=lambda x: os.path.basename(x))

# sorted_image_files = [r"/home/champrin/Desktop/MV-CS016-10UC+DA1041860/red/Video_20250113104416448/exchange_red_Video_20250113104416448_27.jpg"]

# 打印排序后的图片文件路径
for file_path in sorted_image_files:
    print(file_path)

    # Run batched inference on a list of images
    # results = model(sorted_image_files)  # return a list of Results objects
    results = model(file_path ,conf=0.05,iou=0.45)  # return a list of Results objects

    rst = []

    # print(results)

    classes = {'box':0, 'arrow_right':1, 'arrow_left':2}

    def normalize(value, max_value):
        return value / max_value

    # Process results list
    for result in results:
        # 预测框的 xyxy 坐标
        bboxes_xywh = result.boxes.xywh.cpu().numpy().astype('float32')
        bboxes_cls = result.boxes.cls.cpu().numpy().astype('float32')
        bboxes_conf = result.boxes.conf.cpu().numpy().astype('float32')
        # 关键点的 xy 坐标
        bboxes_keypoints = result.keypoints.cpu().data.numpy().astype('float32')

        # print(result)
        # print(result.boxes)
        # print(result.keypoints)
        # print(bboxes_cls)
        # print(bboxes_conf)

        for box_id in range(len(bboxes_xywh)):
            bbox_xyxy = bboxes_xywh[box_id]  # 获取框坐标
            bbox_cls = bboxes_cls[box_id]  # 获取框的预测类别（对于关键点检测，只有一个类别）
            bbox_conf = bboxes_conf[box_id]  # 获取框的预测类别（对于关键点检测，只有一个类别）
            bbox_keypoints = bboxes_keypoints[box_id]  # 该框所有关键点坐标和置信度

            # print(bbox_xyxy)
            # print(bbox_keypoints)
            # box_xy1=(bbox_xyxy[0], bbox_xyxy[1]), 110.0, 47.0, 834.0, 1007.0,
            # box_xy2=(bbox_xyxy[2], bbox_xyxy[3]),
            # keypoints=bbox_keypoints
            rst.append([
                float(bbox_xyxy[0] - bbox_xyxy[2] / 2),
                float(bbox_xyxy[1] - bbox_xyxy[3] / 2),
                float(bbox_xyxy[2]),
                float(bbox_xyxy[3]),
                int(bbox_cls), float(bbox_conf),
                float(bbox_keypoints[0][0]),
                float(bbox_keypoints[0][1]),
                float(bbox_keypoints[0][2]),
                float(bbox_keypoints[1][0]),
                float(bbox_keypoints[1][1]),
                float(bbox_keypoints[1][2]),
                float(bbox_keypoints[2][0]),
                float(bbox_keypoints[2][1]),
                float(bbox_keypoints[2][2]),
                float(bbox_keypoints[3][0]),
                float(bbox_keypoints[3][1]),
                float(bbox_keypoints[3][2]),
                float(bbox_keypoints[4][0]),
                float(bbox_keypoints[4][1]),
                float(bbox_keypoints[4][2]),
                float(bbox_keypoints[5][0]),
                float(bbox_keypoints[5][1]),
                float(bbox_keypoints[5][2]),
                        ])

        # print(rst)


        image = cv2.imread(file_path)
        image_height, image_width, _ = image.shape
        # print((image_height, image_width))
        slots = slot_opt.Detect(image, rst)
        slot_opt.DrawDebug(image)

        yolo_format_data = []
        for slot in slots:
            x, y, w, h = slot.rect

            center_x = x + w / 2
            center_y = y + h / 2
            width = w
            height = h

            # 归一化边界框坐标
            center_x = normalize(center_x, image_width)
            center_y = normalize(center_y, image_height)
            width = normalize(width, image_width)
            height = normalize(height, image_height)

            # 保留6位小数
            center_x = round(center_x, 6)
            center_y = round(center_y, 6)
            width = round(width, 6)
            height = round(height, 6)

            class_id = slot.type.value

            keypoints = []
            cnt = 0
            for keypoint in slot.opt_pts:
                keypoint_x = keypoint[0]
                keypoint_y = keypoint[1]
                keypoint_visibility = 2  # Assuming all keypoints are visible 0：没有明显露出，不可见 1：被遮挡，不可见 2：可见
                keypoint_x = normalize(keypoint_x, image_width)
                keypoint_y = normalize(keypoint_y, image_height)
                # 保留6位小数
                keypoint_x = round(keypoint_x, 6)
                keypoint_y = round(keypoint_y, 6)
                cnt += 1
                keypoints.extend([keypoint_x, keypoint_y, keypoint_visibility])

            if class_id == 0:
                if cnt != 6:
                    for i in range(6 - cnt):
                        keypoints.extend([20., 20., 0])
                try:
                    yolo_format_data.append(f"{class_id} {center_x} {center_y} {width} {height} "
                                            f"{keypoints[0]} {keypoints[1]} {keypoints[2]} "
                                            f"{keypoints[3]} {keypoints[4]} {keypoints[5]} "
                                            f"{keypoints[9]} {keypoints[10]} {keypoints[11]} "
                                            f"{keypoints[15]} {keypoints[16]} {keypoints[17]} "
                                            f"{keypoints[6]} {keypoints[7]} {keypoints[8]} "
                                            f"{keypoints[12]} {keypoints[13]} {keypoints[14]}"
                                            )
                except IndexError as e:
                    print(e)
            else:
                yolo_format_data.append(f"{class_id} {center_x} {center_y} {width} {height} " + " ".join(map(str, keypoints)))

        print(yolo_format_data)

        output_folder = os.path.join(os.path.dirname(result.path), "yolo_annotations")
        os.makedirs(output_folder, exist_ok=True)
        filename = os.path.basename(result.path)
        output_file_path = os.path.join(output_folder, filename.replace(".jpg", ".txt"))

        print(len(yolo_format_data))
        if len(yolo_format_data) != 0:

            with open(output_file_path, 'w') as output_file:
                for line in yolo_format_data:
                    output_file.write(line + "\n")

            # cv2.namedWindow("warp_img", 0)
            # cv2.imshow("warp_img", slot_opt.warp_img)
            # cv2.namedWindow("warp_roi", 0)
            # cv2.imshow("warp_roi", slot_opt.warp_roi)
            # cv2.namedWindow("warp_roib", 0)
            # cv2.imshow("warp_roib", slot_opt.warp_roib)
    #         cv2.namedWindow("image", 0)
    #         cv2.imshow("image", image)
    #         cv2.namedWindow("binary_img_", 0)
    #         cv2.imshow("binary_img_", slot_opt.binary_img_)
    # key = cv2.waitKey(0)
    # if key == 27:
    #     break

cv2.destroyAllWindows()